import pandas as pd
import numpy as np
import time
import datetime
import warnings

warnings.filterwarnings('ignore')
from jqfactor import *
from jqdata import *

import scipy.stats as st
import statsmodels.api as sm
from jqfactor import get_factor_values

import matplotlib as mpl
import matplotlib.pyplot as plt

mpl.rcParams['font.sans-serif'] = ['SimHei']  # 指定默认字体
mpl.rcParams['axes.unicode_minus'] = False  # 解决保存图像是负号'-'显示为方块的问题
import seaborn as sns

sns.set_style({'font.sans-serif': ['simhei', 'Arial']})

# 设定最大显示行数、列数为200
pd.set_option('display.max_rows', 150)
pd.set_option('display.max_columns', 150)
pd.set_option('display.width', 150)

factor_name = [

    '流通市值',
    '总市值',
    '营业收入增长率',
    '总资产增长率',
    '经营活动产生的现金流量净额增长率',
    '利润总额增长率',
    '归属母公司股东的净利润增长率',
    '筹资活动产生的现金流量净额增长率',
    '净利润增长率',
    '净资产增长率',
    '市盈率相对盈利增长比率'
]

factors = [
    'operating_revenue_growth_rate',  # 营业收入增长率
    'total_asset_growth_rate',  # 总资产增长率
    'net_operate_cashflow_growth_rate',  # 经营活动产生的现金流量净额增长率
    'total_profit_growth_rate',  # 利润总额增长率
    'np_parent_company_owners_growth_rate',  # 归属母公司股东的净利润增长率
    'financing_cash_growth_rate',  # 筹资活动产生的现金流量净额增长率
    'net_profit_growth_rate',  # 净利润增长率
    'net_asset_growth_rate',  # 净资产增长率
    'PEG'  # 市盈率相对盈利增长比率
]


def get_monthly_trade_days(start, end):
    """
    获取每月的最后一个交易日。

    :param start: 开始日期
    :param end: 结束日期，应为月末
    :return: 每月最后一个交易日的列表，日期格式为datetime.date
    """
    # 获取交易日历
    trade_days = get_trade_days(start_date=start, end_date=end)
    df = pd.DataFrame(trade_days, index=trade_days)
    # 转换DatetimeIndex
    df.index = pd.to_datetime(df.index)
    # 获取每个月最后一个交易
    monthly_last_trade_days = list(df.resample('M').last().iloc[:, 0])
    return monthly_last_trade_days


# 399303.XSHE


def filter_stocks(trade_date, index_id='399303.XSHE'):
    """
    获取指定日期筛选后的指数成分股。

    :param trade_date: 指定日期
    :param index_id: 指数ID
    :return: 筛选后的股票列表
    """
    # 获取交易日成分股
    stocks = get_index_stocks(index_id, date=trade_date)

    # 获取交易日当天是否处于ST状态
    is_st = get_extras('is_st', stocks, end_date=trade_date, count=1).iloc[0]
    # 获取交易日当天是否全天停牌
    is_suspended = get_price(stocks, end_date=trade_date, count=1, fields='paused', panel=False).set_index('code')[
                       'paused'] == 1

    # 筛选非ST且非停牌的股票
    filtered_stocks = [stock for stock in stocks if not (is_st[stock] or is_suspended[stock])]

    return filtered_stocks


def get_monthly_yield(stock_list, start, end):
    """
    获取月度收益率。

    :param stock_list: 股票列表
    :param start: 开始日期
    :param end: 结束日期
    :return: 股票的月度收益率
    """
    start_prices = get_price(stock_list, count=1, end_date=start, fields='close', panel=False).set_index('code')[
        ['close']]
    end_prices = get_price(stock_list, count=1, end_date=end, fields='close', panel=False).set_index('code')[['close']]
    return_df = pd.concat([start_prices, end_prices], axis=1)
    return_df['return'] = (return_df.iloc[:, 1] - return_df.iloc[:, 0]) / return_df.iloc[:, 0]
    return return_df[['return']]


# def clean_factor(factors,date):
def get_cleaned_factors(factors, date):
    """
    使用聚宽提供的API清洗因子数据：处理缺失值、去极值、标准化、行业市值中性化。

    :param factors: 原始因子数据
    :param date: 数据日期
    :return: 清洗后的因子数据
    """
    # 缺失值处理：使用列均值填充
    factors = factors.fillna(factors.mean())

    # 去极值
    factors = winsorize_med(factors, scale=10, inclusive=True, inf2nan=True, axis=0)

    # 标准化
    factors = standardlize(factors, inf2nan=True, axis=0)

    # 行业市值中性化
    factors = neutralize(factors, ['sw_l1', 'market_cap'], date=str(date), axis=0)

    return factors


def get_all_data(start, end):
    """
    获取指定日期范围内的所有因子数据。

    :param start: 开始日期
    :param end: 结束日期
    :return: 所有截面期的因子数据字典
    """
    all_data_dict = {}
    trade_days = get_monthly_trade_days(start, end)
    print("获取交易日历")
    print(trade_days)

    for idx, date in enumerate(trade_days[:-1]):
        print(f"处理 {date} 的因子数据")

        stock_list = filter_stocks(date)
        financial_factors = get_financial_factors(stock_list, date)
        other_factors_data = get_other_factors(stock_list, date)

        combined_factors = pd.concat([financial_factors, other_factors_data], axis=1)
        # 清洗数据
        cleaned_factors = get_cleaned_factors(combined_factors, date)
        # 合并数据
        next_period_return = get_monthly_yield(stock_list, date, trade_days[idx + 1])
        combined_data = pd.concat([cleaned_factors, next_period_return], axis=1).dropna()
        all_data_dict[str(date)] = combined_data

        print(f"完成处理 {date} 的因子数据")
        print("*" * 30)

    print("所有因子数据获取完毕!")
    return all_data_dict


def get_financial_factors(stock_list, date):
    """
    获取财务指标数据。

    :param stock_list: 股票列表
    :param date: 查询日期
    :return: 财务因子的DataFrame
    """
    query_statement = query(
        valuation.code,
        valuation.circulating_market_cap,  # 流通市值
        valuation.market_cap  # 总市值
    ).filter(valuation.code.in_(stock_list))

    return get_fundamentals(query_statement, date).set_index('code')


def get_other_factors(stock_list, date):
    """
    获取其他因子数据。

    :param stock_list: 股票列表
    :param date: 查询日期
    :return: 其他因子的DataFrame
    """

    other_factors_data = pd.DataFrame(index=stock_list, columns=factors)
    for stock in stock_list:
        stock_data = get_factor_values(stock, factors=factors, end_date=date, count=1)
        for factor, value in stock_data.items():
            other_factors_data.loc[stock][factor] = value.loc[date][stock]

    return other_factors_data


def calculate_ic_beta_t(all_data_dict, factor_names):
    """
    计算所有因子的信息系数(IC)、因子收益率(beta)和t值。

    :param all_data_dict: 所有截面期的因子数据字典
    :param factor_names: 因子名称列表
    :return: 三个DataFrame，分别包含IC、beta和t值
    """
    factors_ic = {}
    factors_beta = {}
    factors_t = {}

    for date, data in all_data_dict.items():
        # 检查数据是否为空
        if data.empty:
            print(f"数据为空: {date}")
            continue

        ic_values = []
        beta_values = []
        t_values = []

        for factor in data.columns[:-1]:

            # 检查每个因子的数据是否足够
            if data[factor].empty or data['return'].empty:
                print(f"因子或收益数据为空: {factor} on {date}")
                continue

            # 计算与最后一列的收益秩相关系数
            ic = st.spearmanr(data[factor], data['return'])[0]
            ic_values.append(ic)

            # 计算beta因子收益率和t值
            # 每列因子与收益率RLM回归,得到系数,t值
            # 加截距,变成二维
            x = sm.add_constant(data[factor])
            # print(x)
            model = sm.RLM(data['return'], x).fit()
            beta_values.append(model.params[1])
            # 计算T
            t_values.append(model.tvalues[1])

        factors_ic[date] = ic_values
        factors_beta[date] = beta_values
        factors_t[date] = t_values
        print(f"{date}数据已处理")

    ic_df = pd.DataFrame(factors_ic, index=factor_names).T
    beta_df = pd.DataFrame(factors_beta, index=factor_names).T
    t_df = pd.DataFrame(factors_t, index=factor_names).T

    return ic_df, beta_df, t_df


def evaluation(ic_df, beta_df, t_df):
    """
    对因子进行综合评估。

    :param ic_df: IC值的DataFrame
    :param beta_df: Beta值的DataFrame
    :param t_df: t值的DataFrame
    :return: 评估结果的DataFrame
    """
    eval_dict = {}

    for factor in ic_df.columns:
        ic_stats = calculate_ic_statistics(ic_df[factor])
        beta_stats = calculate_beta_statistics(beta_df[factor])
        t_stats = calculate_t_statistics(t_df[factor])

        eval_dict[factor] = {**ic_stats, **beta_stats, **t_stats}

    eval_df = pd.DataFrame(eval_dict).T
    eval_df['分数'] = eval_df.apply(calculate_score, axis=1)
    eval_df.sort_values(by='分数', ascending=False, inplace=True)

    return eval_df


def calculate_ic_statistics(ic_series):
    """
    计算IC系列的统计指标。

    :param ic_series: IC值的pandas Series
    :return: 一个包含IC统计指标的字典
    """
    mean = np.round(ic_series.mean(), 4)
    std = np.round(ic_series.std(), 4)
    ratio_002 = np.round((ic_series.abs() > 0.02).mean() * 100, 2)
    ratio_positive = np.round((ic_series > 0).mean() * 100, 2)
    ir = np.round(mean / std, 4)

    return {
        'IC均值': mean,
        'IC标准差': std,
        '|IC|>0.02比率': f"{ratio_002}%",
        'IC>0比率': f"{ratio_positive}%",
        'IR': ir
    }


def calculate_beta_statistics(beta_series):
    """
    计算Beta系列的统计指标。

    :param beta_series: Beta值的pandas Series
    :return: 一个包含Beta统计指标的字典
    """
    mean = np.round(beta_series.mean(), 4)
    t_stat, p_value = np.round(st.ttest_1samp(beta_series, 0), 4)

    return {
        '因子收益率均值': mean,
        '因子收益率t检验': t_stat,
        'p值': p_value
    }


def calculate_t_statistics(t_series):
    """
    计算t值系列的统计指标。

    :param t_series: t值的pandas Series
    :return: 一个包含t值统计指标的字典
    """
    mean_abs_t = np.round(t_series.abs().mean(), 4)
    ratio_greater_2 = np.round((t_series.abs() > 2).mean() * 100, 2)

    return {
        '|t|均值': mean_abs_t,
        '|t|>2比率': f"{ratio_greater_2}%"
    }


def calculate_score(row):
    """
    根据给定的行计算因子的得分。

    :param row: DataFrame的行，包含因子的统计指标
    :return: 计算得到的得分
    """
    score = 0
    if abs(row['IC均值']) > 0.02:
        score += 0.25
    if row['p值'] < 0.05:
        score += 0.25
    if abs(row['IR']) > 0.3:
        score += 0.25
    if row['|t|均值'] > 2:
        score += 0.25
    return score


def plot_heatmap(ic_df, evaluation_df):
    """
    绘制因子间相关性的热力图。

    :param ic_df: 信息系数(IC)的DataFrame
    :param evaluation_df: 因子评估结果的DataFrame
    """
    # 筛选表现良好的因子
    good_factors = evaluation_df.index
    # good_factors = evaluation_df[evaluation_df['分数'] >= 3].index

    selected_ic_df = ic_df[good_factors]
    corr_matrix = selected_ic_df.corr(method='spearman').round(2)
    # 以corr_matrix的形状生成一个全为0的矩阵
    # 绘制热力图
    plt.figure(figsize=(20, 10))
    sns.heatmap(corr_matrix, annot=True, cmap='RdPu')
    plt.show()

    return corr_matrix


def plot_heat(ic_df, eval_df):
    # 截取表现好的因子
    eval_df = eval_df[eval_df['分数'] >= 1]
    ic_df = ic_df.loc[:, eval_df.index]
    corr = np.around(ic_df.corr('spearman'), 2)

    #     #以corr的形状生成一个全为0的矩阵
    #     mask = np.zeros_like(corr)
    #     #将mask的对角线及以上设置为True
    #     #这部分就是对应要被遮掉的部分
    #     mask[np.triu_indices_from(mask)] = True

    fig, ax = plt.subplots(figsize=(20, 10))
    sns.heatmap(corr.abs(), annot=True, cmap='RdPu', ax=ax)
    return corr
